import cv2
import os
import sys
import numpy as np
from PIL import Image
import easyocr

# 全局初始化EasyOCR读取器（只需加载一次模型） 
reader = easyocr.Reader(['ch_sim', 'en'], gpu=False)  # 使用中文简体和英文，不使用GPU 

def detect_colors(image): 
    """自动检测消息框颜色配置""" 
    # 添加详细日志记录
    print("开始颜色检测...")
    print(f"图像尺寸: {image.shape}")

    # 使用K-means聚类找到主要颜色
    print("执行K-means聚类...")
    pixels = image.reshape(-1, 3).astype(np.float32)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    _, labels, centers = cv2.kmeans(pixels, 4, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    print(f"检测到的颜色中心: {centers}")
    
    # 分析颜色特征
    color_profile = {
        "user_bubble": None,
        "other_bubble": None,
        "text_color": None
    }
    
    # 根据HSV值识别消息框颜色
    hsv_centers = [cv2.cvtColor(np.uint8([[c]]), cv2.COLOR_BGR2HSV)[0][0] for c in centers]
    for i, hsv in enumerate(hsv_centers):
        print(f"分析颜色 {i+1} (HSV): {hsv}")
        # 识别用户消息框（高饱和度绿色）- 扩大色域范围
        if 30 < hsv[0] < 90 and hsv[1] > 60:
            color_profile["user_bubble"] = tuple(map(int, centers[i]))
            print(f"识别为用户消息框颜色: {centers[i]}")
        # 识别他人消息框（低饱和度白色）
        elif hsv[1] < 50 and hsv[2] > 200:
            color_profile["other_bubble"] = tuple(map(int, centers[i]))
            print(f"识别为他人消息框颜色: {centers[i]}")
        # 识别文本颜色（低亮度）
        elif hsv[2] < 100:
            color_profile["text_color"] = tuple(map(int, centers[i]))
            print(f"识别为文本颜色: {centers[i]}")
    
    print("颜色检测完成，配置如下:")
    print(color_profile)
    return color_profile

def detect_bubbles(image, color_profile=None):
    """自动检测消息框区域"""
    if color_profile is None:
        color_profile = detect_colors(image)
    
    # 动态调整颜色阈值
    masks = {
        "user": cv2.inRange(image, 
            np.array(color_profile["user_bubble"]) - 50 if color_profile["user_bubble"] else np.array([0, 0, 0]),
            np.array(color_profile["user_bubble"]) + 50 if color_profile["user_bubble"] else np.array([0, 0, 0])),
        "other": cv2.inRange(image,
            np.array(color_profile["other_bubble"]) - 20 if color_profile["other_bubble"] else np.array([0, 0, 0]),
            np.array(color_profile["other_bubble"]) + 20 if color_profile["other_bubble"] else np.array([0, 0, 0]))
    }
    
    # 查找轮廓
    bubbles = {}
    for speaker, mask in masks.items():
        # 形态学操作改善轮廓
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        # 查找轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 过滤小轮廓
        min_area = image.shape[0] * image.shape[1] * 0.005  # 最小面积阈值
        valid_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > min_area]
        
        # 保存轮廓和边界框
        bubbles[speaker] = []
        for cnt in valid_contours:
            x, y, w, h = cv2.boundingRect(cnt)
            bubbles[speaker].append({
                'contour': cnt,
                'bbox': (x, y, w, h),
                'area': cv2.contourArea(cnt)
            })
    
    return bubbles

def draw_annotations(image, bubbles, color_profile, output_path):
    """在图像上标注识别的对话框和颜色，并保存结果"""
    print(f"开始标注对话框和颜色...")
    
    # 创建图像副本，以便在其上绘制而不影响原图
    annotated_img = image.copy()
    
    # 为每种对话框类型定义不同的边框颜色（BGR格式）
    border_colors = {
        'user': (0, 0, 255),  # 红色
        'other': (255, 0, 0)  # 蓝色
    }
    
    # 标注每个对话框
    bubble_count = 0
    for speaker, bubble_list in bubbles.items():
        for bubble in bubble_list:
            x, y, w, h = bubble['bbox']
            
            # 绘制矩形框
            cv2.rectangle(annotated_img, (x, y), (x+w, y+h), border_colors[speaker], 2)
            
            # 添加标签（说话者和区域信息）
            label = f"{speaker}: Area={int(bubble['area'])}"
            cv2.putText(annotated_img, label, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, border_colors[speaker], 2)
            
            bubble_count += 1
    
    # 绘制颜色配置信息
    # 在图像底部添加颜色信息
    info_y = image.shape[0] - 10  # 底部位置

    if color_profile['user_bubble']:
        user_color = color_profile['user_bubble']
        cv2.putText(annotated_img, f"用户对话框颜色: {user_color}", 
                   (10, info_y - 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        # 绘制颜色示例框
        cv2.rectangle(annotated_img, (10, info_y - 50), (50, info_y - 20), user_color, -1)
    
    if color_profile['other_bubble']:
        other_color = color_profile['other_bubble']
        cv2.putText(annotated_img, f"他人对话框颜色: {other_color}", 
                   (10, info_y - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
        # 绘制颜色示例框
        cv2.rectangle(annotated_img, (10, info_y - 20), (50, info_y), other_color, -1)
    
    # 创建输出目录（如果不存在）
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 保存带标注的图像
    cv2.imwrite(output_path, annotated_img)
    print(f"已保存标注图像到: {output_path}")
    print(f"共标注了 {bubble_count} 个对话框")
    
    return annotated_img

def extract_text(image, bbox):
    """从消息框中提取文本 - 使用EasyOCR替代Tesseract"""
    x, y, w, h = bbox
    roi = image[y:y+h, x:x+w]
    
    # 使用EasyOCR直接识别图像
    results = reader.readtext(roi)
    
    # 合并所有识别结果
    if results:
        # 提取所有文本并按垂直位置排序
        texts = [text for _, text, conf in results if conf > 0.2]  # 过滤低置信度结果
        return ' '.join(texts)
    
    return ""

def analyze_background_color(image, bbox, padding=5):
    """分析文本区域的背景颜色，判断消息类型"""
    x, y, w, h = bbox
    
    # 扩大区域以包含更多背景
    x_expanded = max(0, x - padding)
    y_expanded = max(0, y - padding)
    w_expanded = min(image.shape[1] - x_expanded, w + 2 * padding)
    h_expanded = min(image.shape[0] - y_expanded, h + 2 * padding)
    
    # 提取扩展区域
    roi = image[y_expanded:y_expanded+h_expanded, x_expanded:x_expanded+w_expanded]
    
    # 创建文本区域掩码（这只是一个粗略的掩码，实际应用中可能需要更精确的文字分割）
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    _, text_mask = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY_INV)
    
    # 创建背景掩码 (排除文本)
    background_mask = cv2.bitwise_not(text_mask)
    
    # 提取背景区域
    background = cv2.bitwise_and(roi, roi, mask=background_mask)
    
    # 计算背景区域的平均颜色（排除黑色像素）
    non_zero = background[np.where((background > [0,0,0]).all(axis=2))]    
    if len(non_zero) > 0:
        avg_color = np.mean(non_zero, axis=0).astype(int)
    else:
        return "unknown"  # 如果没有足够的背景像素
    
    # 转换为HSV进行颜色分析
    hsv_color = cv2.cvtColor(np.uint8([[avg_color]]), cv2.COLOR_BGR2HSV)[0][0]
    
    # 根据HSV值判断是用户消息（绿色）还是他人消息（白色）或其他
    h, s, v = hsv_color
    
    # 判断条件
    if 30 < h < 90 and s > 60:  # 绿色(用户消息)
        return "user"
    elif s < 50 and v > 200:  # 白色(他人消息)
        return "other"
    elif 70 < v < 200 and s < 30:  # 灰色(引用内容)
        return "quote"
    else:
        return "unknown"

def should_filter_text(text, msg_type):
    """判断文本是否应该被过滤掉"""
    # 过滤掉引用内容（灰色框内的文字）
    if msg_type == "quote":
        return True
        
    # 过滤手机状态栏相关内容
    status_patterns = [
        r'\d+:\d+\s*(上午|下午|AM|PM)?',  # 手机时间，如 3:03，下午2:57
        r'\d+%',                       # 电池百分比
        r'[2-5]G',                      # 网络信号，如5G, 4G
        r'Wi-?Fi',                     # WiFi标识
        r'微信',                        # 微信标识
        r'返回',                        # 返回按钮
        r'举报',                        # 举报按钮
        r'翻译',                        # 翻译选项
        r'复制',                        # 复制选项
        r'删除',                        # 删除选项
        r'转发',                        # 转发选项
        r'(上午|下午)\s*\d+:\d+',       # 对话时间，如 下午2:57
        r'昨天|前天|今天|星期[一二三四五六日天]', # 日期标记
        r'群聊|聊天室',                  # 群聊标识
        r'\d+条未读',                   # 未读消息数
        r'置顶|免打扰',                  # 对话设置选项
    ]
    
    # 检查是否匹配任何模式
    import re
    for pattern in status_patterns:
        if re.search(pattern, text):
            # 如果整个文本都是模式，或者文本很短，则过滤整个文本
            if len(text) < 15 or re.match(f'^\\s*{pattern}\\s*$', text):
                return True
    
    return False

def process_image_by_text_first(image_path, output_dir=None):
    """使用先文本识别后背景色分析的新方法处理图像"""
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"无法读取图像: {image_path}")
    
    print(f"图片尺寸: {image.shape[1]}x{image.shape[0]}")
    print("使用新方法: 先OCR识别文本，后分析背景色")
    
    # 使用EasyOCR直接识别图像中所有文本
    print("开始OCR识别文本...")
    results = reader.readtext(image)
    print(f"识别到 {len(results)} 个文本区域")
    
    # 分析每个文本区域的背景色并分类
    dialogues = []
    position = 0
    
    # 存储所有识别结果，用于标注和显示
    text_areas = []
    
    for i, (bbox, text, conf) in enumerate(results):
        if conf < 0.2:  # 过滤低置信度结果
            continue
            
        # 转换bbox格式为(x,y,w,h)
        tl, tr, br, bl = bbox
        x = int(min(tl[0], bl[0]))
        y = int(min(tl[1], tr[1]))
        w = int(max(tr[0], br[0]) - x)
        h = int(max(bl[1], br[1]) - y)
        
        bbox_rect = (x, y, w, h)
        
        # 分析文本区域背景色判断消息类型
        msg_type = analyze_background_color(image, bbox_rect)
        
        # 过滤掉应该过滤的文本
        if should_filter_text(text, msg_type):
            continue
        
        # 保存分析结果
        text_areas.append({
            'bbox': bbox_rect,
            'points': bbox,  # 原始四边形点
            'text': text,
            'type': msg_type,
            'y_center': y + h // 2  # 用于排序
        })
        
    # 按垂直位置排序
    text_areas.sort(key=lambda x: x['y_center'])
    
    # 将相邻且类型相同的文本区域合并为一条消息
    messages = []
    current_msg = None
    
    for area in text_areas:
        if current_msg is None or current_msg['type'] != area['type'] or \
           area['y_center'] - current_msg['y_end'] > 50:  # 如果距离太远，认为是新消息
            if current_msg is not None:
                messages.append(current_msg)
            current_msg = {
                'type': area['type'],
                'text': area['text'],
                'y_start': area['y_center'],
                'y_end': area['y_center'],
                'areas': [area]
            }
        else:
            # 属于同一消息，合并文本
            current_msg['text'] += " " + area['text']
            current_msg['y_end'] = area['y_center']
            current_msg['areas'].append(area)
    
    if current_msg is not None:
        messages.append(current_msg)
    
    # 转换为对话格式
    for i, msg in enumerate(messages):
        # 跳过引用内容
        if msg['type'] == 'quote':
            continue
            
        if msg['type'] in ['user', 'other']:
            speaker = '用户' if msg['type'] == 'user' else '他人'
            dialogues.append({
                'position': i,
                'speaker': speaker,
                'content': msg['text']
            })
    
    # 如果提供了输出目录，生成标注图像
    if output_dir:
        # 创建输出文件路径
        base_name = os.path.basename(image_path)
        output_path = os.path.join(output_dir, f"text_based_{base_name}")
        # 绘制标注
        draw_text_annotations(image, text_areas, output_path)
    
    print(f"发现 {len(dialogues)} 条消息...")
    return dialogues

def draw_text_annotations(image, text_areas, output_path):
    """在图像上标注识别的文本区域和消息类型"""
    print("开始标注文本区域和背景颜色...")
    
    # 创建图像副本
    annotated_img = image.copy()
    
    # 颜色映射
    color_map = {
        'user': (0, 255, 0),    # 绿色
        'other': (0, 0, 255),   # 红色
        'quote': (128, 128, 128), # 灰色 - 这些将会被过滤掉，但保留颜色便于调试
        'unknown': (255, 255, 0)  # 黄色
    }
    
    # 在图像顶部添加过滤说明
    filter_info = "已过滤: 手机状态栏信息、对话时间、引用内容"
    cv2.putText(annotated_img, filter_info, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    
    # 标注每个文本区域
    for i, area in enumerate(text_areas):
        # 绘制文本区域边界框
        x, y, w, h = area['bbox']
        color = color_map.get(area['type'], (255, 255, 0))  # 默认黄色
        cv2.rectangle(annotated_img, (x, y), (x+w, y+h), color, 2)
        
        # 添加标签（类型和文本）
        label = f"{area['type']}: {area['text'][:10]}{'...' if len(area['text']) > 10 else ''}"
        cv2.putText(annotated_img, label, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
    
    # 创建输出目录（如果不存在）
    output_dir = os.path.dirname(output_path)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 保存带标注的图像
    cv2.imwrite(output_path, annotated_img)
    print(f"已保存文本标注图像到: {output_path}")
    print(f"共标注了 {len(text_areas)} 个文本区域")
    
    return annotated_img

def process_image(image_path, output_dir=None, use_text_first=True):
    """处理微信对话截图并提取对话内容"""
    # 使用新方法还是旧方法
    if use_text_first:
        return process_image_by_text_first(image_path, output_dir)
    
    # 以下为旧方法
    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"无法读取图像: {image_path}")
    
    print(f"图片尺寸: {image.shape[1]}x{image.shape[0]}")
    
    # 检测颜色配置
    color_profile = detect_colors(image)
    
    # 检测消息框
    bubbles = detect_bubbles(image, color_profile)
    
    # 如果提供了输出目录，生成标注图像
    if output_dir:
        # 创建输出文件路径
        base_name = os.path.basename(image_path)
        output_path = os.path.join(output_dir, f"annotated_{base_name}")
        # 绘制标注
        draw_annotations(image, bubbles, color_profile, output_path)
    
    # 合并所有消息框
    all_bubbles = []
    for speaker, bubble_list in bubbles.items():
        for bubble in bubble_list:
            all_bubbles.append({
                'speaker': '用户' if speaker == 'user' else '他人',
                'bbox': bubble['bbox'],
                'area': bubble['area'],
                'y_center': bubble['bbox'][1] + bubble['bbox'][3] // 2
            })
    
    # 按垂直位置排序
    all_bubbles.sort(key=lambda x: x['y_center'])
    
    # 提取文本
    dialogues = []
    for i, bubble in enumerate(all_bubbles):
        text = extract_text(image, bubble['bbox'])
        if text:  # 只添加非空文本
            dialogues.append({
                'position': i,
                'speaker': bubble['speaker'],
                'content': text
            })
    
    print(f"发现{len(dialogues)}个消息框...")
    return dialogues

def batch_test(folder_path, use_text_first=True):
    """批量测试文件夹中的图片"""
    if not os.path.exists(folder_path):
        print(f"错误：文件夹 '{folder_path}' 不存在")
        return

    valid_files = [
        f for f in os.listdir(folder_path) 
        if f.lower().endswith(('.png', '.jpg', '.jpeg'))
        and os.path.isfile(os.path.join(folder_path, f))
    ]
    
    if not valid_files:
        print(f"⚠ 警告：文件夹 '{folder_path}' 中没有找到图片文件")
        return
    
    # 创建输出目录
    method_name = "text_first" if use_text_first else "bubble_first"
    output_dir = os.path.join(folder_path, f"annotated_{method_name}")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建标注图像输出目录: {output_dir}")

    method_desc = "先识别文本后分析背景色" if use_text_first else "先检测气泡后提取文本"
    print(f"\n使用方法: {method_desc}")
    print(f"正在处理文件夹中的 {len(valid_files)} 张图片：")
    
    for filename in valid_files:
        full_path = os.path.join(folder_path, filename)
        print(f"\n▷ 当前文件: {filename}")
        
        try:
            # 处理图像并生成标注
            dialogues = process_image(full_path, output_dir, use_text_first)
            if dialogues:
                print("✅ 成功提取内容：")
                for d in dialogues:
                    print(f" {d['position']+1}. {d['speaker']}: {d['content']}")
            else:
                print("⚠ 未提取到有效对话")
        except Exception as e:
            print(f"❌ 处理失败：{str(e)}")

def main():
    # 默认使用文本优先方法
    use_text_first = True
    
    if len(sys.argv) > 1 and sys.argv[1] == "--bubble-first":
        use_text_first = False
        # 移除选项参数
        sys.argv.pop(1)
    
    if len(sys.argv) == 2 and os.path.isdir(sys.argv[1]):
        batch_test(sys.argv[1], use_text_first)
    elif len(sys.argv) == 2:
        # 为单个文件创建output目录
        method_name = "text_first" if use_text_first else "bubble_first"
        output_dir = os.path.join(os.path.dirname(sys.argv[1]), f"annotated_{method_name}")
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"创建标注图像输出目录: {output_dir}")
        
        method_desc = "先识别文本后分析背景色" if use_text_first else "先检测气泡后提取文本"
        print(f"使用方法: {method_desc}")
        
        dialogues = process_image(sys.argv[1], output_dir, use_text_first)
        print("\n提取结果：")
        for d in dialogues:
            print(f"{d['position']+1}. {d['speaker']}: {d['content']}")
    else:
        print("使用方法：")
        print("单个文件: python wechat_dialog_extractor.py [选项] <图片路径>")
        print("批量测试: python wechat_dialog_extractor.py [选项] <文件夹路径>")
        print("\n可用选项:")
        print("  --bubble-first  使用旧的气泡检测优先方法")

if __name__ == "__main__":
    main()
